import numpy as np
import geatpy as ea

class MyProb(ea.Problem):
    def __init__(self):
        ea.Problem.__init__(
            self,
            name='my_prob',
            M = 2,
            maxormins=[1,1],
            Dim=2,
            varTypes=[0,0],
            lb=[0,0],
            ub=[5,3],
            lbin=[1,1],
            ubin=[1,1]
        )

    def aimFunc(self, pop):
        Vars = pop.Phen
        x1 = Vars[:,[0]]
        x2 = Vars[:,[1]]
        pop.CV = np.hstack([(x1-5)**2 + x2**2 - 25,
                            7.7 - (x1-8)**2 - (x2-3)**2])
        pop.ObjV = np.hstack([4*x1**2 + 4*x2**2,
                              4*(x1-5)**2 + 4*(x2-5)**2])

    def calReferObjV(self):
        N = 10000
        x1 = np.linspace(0,5,N)
        x2 = np.linspace(0,3,N)
        return np.vstack([4*x1**2 + 4*x2**2,
                          4*(x1-5)**2 + 4*(x2-5)**2]).T


if __name__ == '__main__':
    # 建立问题
    myProb = MyProb()
    # 实例化种群
    Field = ea.crtfld(Encoding='RI', varTypes=myProb.varTypes, ranges=myProb.ranges, borders=myProb.borders)
    population = ea.Population(Encoding='RI', Field=Field, NIND=100)
    # 选择算法模板并设置相关参数
    myAlg = ea.moea_NSGA2_templet(problem=myProb, population=population)
    # 算法类统一参数
    myAlg.MAXGEN = 200
    myAlg.logTras = 1
    myAlg.drawing = 1
    myAlg.verbose = False
    # NSAG-2参数，可以自己设置也可以使用默认的
    myAlg.mutOper.Pm = 0.2  # 默认为1
    myAlg.recOper.XOVR = 0.9  # 默认为1
    # 运行
    [NDSet, population] = myAlg.run()  # NDset是整个程序中非支配个体组成的种群，有population类所有属性
    # 输出结果
    print(f'程序运行时间：{myAlg.passTime}s')
    if NDSet.sizes == 0:
        print(f'没有找到可行解！')
    else:
        NDSet.save()
    if myAlg.log is not None and NDSet.sizes != 0:
        print('GD', myAlg.log['gd'][-1])
        print('IGD', myAlg.log['igd'][-1])
        print('HV', myAlg.log['hv'][-1])
        print('Spacing', myAlg.log['spacing'][-1])
    metricName = [['igd'], ['hv']]
    Metrics = np.array([myAlg.log[metricName[i][0]] for i in
                        range(len(metricName))]).T
    # 绘制指标追踪分析图
    ea.trcplot(Metrics, labels=metricName, titles=metricName)
